#include <linux/init.h>  //_initdata
#include <linux/kallsyms.h>
#include <linux/lsm_hooks.h>
#include <linux/module.h>
#include <linux/types.h>  // umode_t
#include <linux/version.h>

#define MAX_RO_PAGES 1024
static struct page*              ro_pages[MAX_RO_PAGES] __initdata;
static unsigned int ro_pages_len __initdata;

struct security_hook_heads probe_dummy_security_hook_heads;

#define MY_HOOK_INIT(HEAD, HOOK)                                               \
    {                                                                          \
        .head = &probe_dummy_security_hook_heads.HEAD, .hook = {.HEAD = HOOK } \
    }

/**
 * probe_security_bprm_committed_creds - Dummy function which does identical to
 * security_bprm_committed_creds() in security/security.c.
 *
 * @bprm: Pointer to "struct linux_binprm".
 *
 * Returns nothing.
 */
void probe_security_bprm_committed_creds(struct linux_binprm* bprm)
{
    do {
        struct security_hook_list* p;

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 17, 0)
        hlist_for_each_entry(
            p, &probe_dummy_security_hook_heads.bprm_committed_creds, list)
            p->hook.bprm_committed_creds(bprm);
#else
        list_for_each_entry(
            p, &probe_dummy_security_hook_heads.bprm_committed_creds, list)
            p->hook.bprm_committed_creds(bprm);
#endif
    } while (0);
}

/**
 * probe_find_symbol - Find function's address from /proc/kallsyms .
 *
 * @keyline: Function to find.
 *
 * Returns address of specified function on success, NULL otherwise.
 */
static void* __init probe_find_symbol(const char* keyline)
{
#if LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 33) && \
    LINUX_VERSION_CODE < KERNEL_VERSION(5, 7, 0)
    char  buf[128] = {};
    char* cp;

    strncpy(buf, keyline + 1, sizeof(buf) - 1);
    cp = strchr(buf, '\n');
    if (cp) *cp = '\0';
    return (void*)kallsyms_lookup_name(buf);
#else
    struct file* file = NULL;
    char* buf;
    unsigned long entry = 0;

    {
#if LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 18)
        struct file_system_type* fstype = get_fs_type("proc");
        struct vfsmount* mnt = vfs_kern_mount(fstype, 0, "proc", NULL);
#elif LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 8)
        struct file_system_type* fstype = NULL;
        struct vfsmount*         mnt = do_kern_mount("proc", 0, "proc", NULL);
#else
        struct file_system_type* fstype = get_fs_type("proc");
        struct vfsmount*         mnt    = kern_mount(fstype);
#endif
        struct dentry* root;
        struct dentry* dentry;

        /*
         * We embed put_filesystem() here because it is not exported.
         */
        if (fstype) module_put(fstype->owner);
        if (IS_ERR(mnt)) goto out;
        root = dget(mnt->mnt_root);
#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 5, 0)
        inode_lock(root->d_inode);
        dentry = lookup_one_len("kallsyms", root, 8);
        inode_unlock(root->d_inode);
#elif LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 16)
        mutex_lock(&root->d_inode->i_mutex);
        dentry = lookup_one_len("kallsyms", root, 8);
        mutex_unlock(&root->d_inode->i_mutex);
#else
        down(&root->d_inode->i_sem);
        dentry = lookup_one_len("kallsyms", root, 8);
        up(&root->d_inode->i_sem);
#endif
        dput(root);
        if (IS_ERR(dentry))
            mntput(mnt);
        else {
#if LINUX_VERSION_CODE >= KERNEL_VERSION(3, 6, 0)
            struct path path = {.mnt = mnt, .dentry = dentry};

            file = dentry_open(&path, O_RDONLY, current_cred());
#else
            file = dentry_open(dentry, mnt, O_RDONLY
#if LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 29)
                               ,
                               current_cred()
#endif
            );
#endif
        }
    }
    if (IS_ERR(file) || !file) goto out;
    buf = kmalloc(PAGE_SIZE, GFP_KERNEL);
    if (buf) {
        int len;
        int offset = 0;

        while ((len = my_kernel_read(file, offset, buf, PAGE_SIZE - 1)) > 0) {
            char* cp;

            buf[len] = '\0';
            cp = strrchr(buf, '\n');
            if (!cp) break;
            *(cp + 1) = '\0';
            offset += strlen(buf);
            cp = strstr(buf, keyline);
            if (!cp) continue;
            *cp = '\0';
            while (cp > buf && *(cp - 1) != '\n') cp--;
            entry = simple_strtoul(cp, NULL, 16);
            break;
        }
        kfree(buf);
    }
    filp_close(file, NULL);
out:
    return (void*)entry;
#endif
}

/**
 * probe_find_variable - Find variable's address using dummy.
 *
 * @function: Pointer to dummy function's entry point.
 * @addr:     Address of the variable which is used within @function.
 * @symbol:   Name of symbol to resolve.
 *
 * This trick depends on below assumptions.
 *
 * (1) @addr is found within 128 bytes from @function, even if additional
 *     code (e.g. debug symbols) is added.
 * (2) It is safe to read 128 bytes from @function.
 * (3) @addr != Byte code except @addr.
 */
static void* __init probe_find_variable(void* function, unsigned long addr,
                                        const char* symbol)
{
    int i;
    u8* base;
    u8* cp = function;

#if LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 24) || \
    LINUX_VERSION_CODE < KERNEL_VERSION(2, 6, 3)
    if (*symbol == ' ')
        base = probe_find_symbol(symbol);
    else
#endif
        base = __symbol_get(symbol);
    if (!base) return NULL;
#if defined(CONFIG_ARM64) && defined(LSM_HOOK_INIT)
    if (function == probe_security_bprm_committed_creds)
        return probe_security_hook_heads_on_arm64((unsigned int*)base);
#endif
#if defined(CONFIG_ARM) && defined(LSM_HOOK_INIT)
    if (function == probe_security_bprm_committed_creds)
        return probe_security_hook_heads_on_arm((unsigned int*)base);
#endif
#if defined(CONFIG_ARM) && LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 24) && \
    !defined(LSM_HOOK_INIT)
    if (function == probe_security_file_alloc)
        return probe_security_ops_on_arm((unsigned int*)base);
#endif
#if defined(CONFIG_ARM) && LINUX_VERSION_CODE < KERNEL_VERSION(2, 6, 36)
    return probe_find_vfsmount_lock_on_arm(function, addr, (unsigned int*)base);
#endif
    /* First, assume absolute adressing mode is used. */
    for (i = 0; i < 128; i++) {
        if (*(unsigned long*)cp == addr) return base + i;
        cp++;
    }
    /* Next, assume PC-relative addressing mode is used. */
#if defined(CONFIG_S390)
    cp = function;
    for (i = 0; i < 128; i++) {
        if ((unsigned long)(cp + (*(int*)cp) * 2 - 2) == addr) {
            static void* cp4ret;

            cp = base + i;
            cp += (*(int*)cp) * 2 - 2;
            cp4ret = cp;
            return &cp4ret;
        }
        cp++;
    }
#endif
    cp = function;
    for (i = 0; i < 128; i++) {
        if ((unsigned long)(cp + sizeof(int) + *(int*)cp) == addr) {
            static void* cp4ret;

            cp = base + i;
            cp += sizeof(int) + *(int*)cp;
            cp4ret = cp;
            return &cp4ret;
        }
        cp++;
    }
    cp = function;
    for (i = 0; i < 128; i++) {
        if ((unsigned long)(long)(*(int*)cp) == addr) {
            static void* cp4ret;

            cp     = base + i;
            cp     = (void*)(long)(*(int*)cp);
            cp4ret = cp;
            return &cp4ret;
        }
        cp++;
    }
    return NULL;
}

static void* __init check_function_address(void* ptr, char* symbol)
{
#if LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 26)
    static char buf[KSYM_SYMBOL_LEN];
    const int   len = strlen(symbol);
#endif

    if (!ptr) {
        printk(KERN_ERR "Can't resolve %s().\n", symbol);
        return NULL;
    }
#if LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 26)
    snprintf(buf, sizeof(buf), "%pS", ptr);
    if (strncmp(buf, symbol, len) || strncmp(buf + len, "+0x0/", 5)) {
        printk(KERN_ERR "Guessed %s is %s\n", symbol, buf);
        return NULL;
    }
#endif
    return ptr;
}

/**
 * probe_security_hook_heads - Find address of "struct security_hook_heads
 * security_hook_heads".
 *
 * Returns pointer to "struct security_hook_heads" on success, NULL otherwise.
 */
struct security_hook_heads* __init probe_security_hook_heads(void)
{
    const unsigned int offset =
        offsetof(struct security_hook_heads, bprm_committed_creds);
    void* cp;

    struct security_hook_heads* shh;
    struct security_hook_list*  entry;
    void* cap = probe_find_symbol(" cap_bprm_set_creds\n");

    /* Get location of cap_bprm_set_creds(). */
    cap = check_function_address(cap, "cap_bprm_set_creds");
    if (!cap) return NULL;
    /* Guess "struct security_hook_heads security_hook_heads;". */
    cp = probe_find_variable(
        probe_security_bprm_committed_creds,
        ((unsigned long)&probe_dummy_security_hook_heads) + offset,
        " security_bprm_committed_creds\n");
    if (!cp) {
        printk(KERN_ERR "Can't resolve security_bprm_committed_creds().\n");
        return NULL;
    }

    /* This should be "struct security_hook_heads security_hook_heads;". */
    shh = ((void*)(*(unsigned long*)cp)) - offset;

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 17, 0)
    hlist_for_each_entry(
        entry, &shh->bprm_set_creds,
        list) if (entry->hook.bprm_set_creds == cap) return shh;
#else
    list_for_each_entry(entry, &shh->bprm_set_creds,
                        list) if (entry->hook.bprm_set_creds == cap) return shh;
#endif
    printk(KERN_ERR "Guessed security_hook_heads is 0x%lx\n",
           (unsigned long)shh);
    return NULL;
}

// int (*inode_mkdir)(struct inode *dir, struct dentry *dentry, umode_t mode);
int linx_inode_mkdir(struct inode* dir, struct dentry* dentry, umode_t mode)
{
    printk("run linx_inode_mkdir\n");
    return 0;
}

// int (*inode_rmdir)(struct inode *dir, struct dentry *dentry);
int linx_inode_rmdir(struct inode* dir, struct dentry* dentry)
{
    printk("run linx_inode_rmdir\n");
    return 0;
}

static struct security_hook_list linx_hooks[] = {
    MY_HOOK_INIT(inode_mkdir, linx_inode_mkdir),
    MY_HOOK_INIT(inode_rmdir, linx_inode_rmdir),
};

static bool __init lsm_test_page_ro(void* addr)
{
    unsigned int i;
    int          unused;
    struct page* page;

    page = (struct page*)lookup_address((unsigned long)addr, &unused);
    if (!page) return 0;
    if (test_bit(_PAGE_BIT_RW, &(page->flags))) return 1;
    for (i = 0; i < ro_pages_len; i++)
        if (page == ro_pages[i]) return 1;
    if (ro_pages_len == MAX_RO_PAGES) return 0;
    ro_pages[ro_pages_len++] = page;
    return 1;
}

static int __init check_ro_pages(struct security_hook_heads* hooks)
{
    int i;

    struct hlist_head* list = &hooks->capable;

    if (!probe_kernel_write(list, list, sizeof(void*))) return 1;

    for (i = 0; i < ARRAY_SIZE(linx_hooks); i++) {
        struct hlist_head*         head = linx_hooks[i].head;
        struct security_hook_list* shp;

        if (!lsm_test_page_ro(&head->first)) return 0;
        hlist_for_each_entry(
            shp, head, list) if (!lsm_test_page_ro(&shp->list.next) ||
                                 !lsm_test_page_ro(&shp->list.pprev)) return 0;
    }
    return 1;
}

static inline void add_hook(struct security_hook_list* hook)
{
    hlist_add_tail_rcu(&hook->list, hook->head);
}

static int __init proc_init(void)
{
    int                         idx;
    struct security_hook_heads* hooks = probe_security_hook_heads();
    if (!hooks) {
        printk("hooks search failed\n");
    }
    for (idx = 0; idx < ARRAY_SIZE(linx_hooks); idx++)
        linx_hooks[idx].head =
            ((void*)hooks) + ((unsigned long)linx_hooks[idx].head) -
            ((unsigned long)&probe_dummy_security_hook_heads);

    if (!check_ro_pages(hooks)) {
        printk(KERN_INFO
               "Can't update security_hook_heads due to write protected. Retry "
               "with rodata=0 kernel command line option added.\n");
        return -EINVAL;
    }

    for (idx = 0; idx < ro_pages_len; idx++)
        set_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));

    for (idx = 0; idx < ARRAY_SIZE(linx_hooks); idx++)
        add_hook(&linx_hooks[idx]);

    for (idx = 0; idx < ro_pages_len; idx++)
        clear_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));

    printk("hook-test install\n");
    return 0;
}

static void proc_cleanup(void)
{
    // security_delete_hooks(my_hook_list, ARRAY_SIZE(my_hook_list));
    printk("hook-test uninstall \n");
}

module_init(proc_init);
module_exit(proc_cleanup);
MODULE_LICENSE("GPL");